#encoding=utf8
import numpy as np
import pickle
import os
import sys
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.svm import LinearSVC
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

# 设置Python的递归深度以避免可能的Stack Overflow错误
sys.setrecursionlimit(2000)

def load_dataset(file_name):
    '''
    从文件读入数据集
    被多处调用，请勿删除或改动本函数！！！
    '''
    try:
        with open(file_name, 'rb') as f:
            raw_dataset = pickle.load(f)
    except FileNotFoundError:
        print(f"错误: 文件 {file_name} 未找到。请确保文件路径正确。")
        return None, None
        
    try:
        example_image = raw_dataset[0][0]
    except KeyError:
        print("错误: 数据集格式不正确，无法找到类别0的数据。")
        return None, None
    except TypeError:
        print("错误: 数据集格式不正确，类别0的数据不是列表或数组。")
        return None, None

    dataset = np.empty((0, example_image.size))
    labels = np.empty((0, 1))

    total_images = 0
    for i_class in raw_dataset.keys():
        images_list = raw_dataset.get(i_class, [])
        if not isinstance(images_list, list) or len(images_list) == 0:
            continue
        for image in images_list:
            features = image.flatten() / 255.0
            dataset = np.vstack((dataset, features))
            labels = np.vstack((labels, i_class))
            total_images += 1
            
    print(f"成功加载 {total_images} 张手写数字图片。")
    return dataset, labels

class Classifier:
    def __init__(self, train_dataset, train_labels):
        self.model = None
        self.pipeline = None
        self.train_dataset = train_dataset
        self.train_labels = train_labels

    def train(self):
        print("=== 正在训练SVM模型... ===")
        # 使用管道串联标准化和SVM模型，因为SVM对特征尺度敏感
        self.pipeline = Pipeline([
            ('scaler', StandardScaler()),
            ('svm', LinearSVC(C=1.0, random_state=42, dual=True, max_iter=2000))
        ])
        
        # 训练模型
        train_labels = self.train_labels.ravel()
        self.pipeline.fit(self.train_dataset, train_labels)
        print("=== 模型训练完成 ===")

    def predict(self, test_dataset):
        '''
        输入：测试数据 test_dataset: 形状为(500, 784)的ndarray
        输出：预测结果 predicted_labels: 形状为(500, )的ndarray
        '''
        predicted_labels = self.pipeline.predict(test_dataset)
        return predicted_labels

def calculate_accuracy(file_name, classifier):
    test_dataset, test_labels = load_dataset(file_name)
    if test_dataset is None:
        return 0
    random_indices = np.random.permutation(test_dataset.shape[0])
    test_dataset = test_dataset[random_indices,:]
    test_labels = test_labels[random_indices,:]
    predicted_labels = classifier.predict(test_dataset)
    if isinstance(predicted_labels, np.ndarray):
        if predicted_labels.size != test_labels.size:
            print('错误：输出的标签数量与测试集大小不一致')
            accuracy = 0
        else:
            accuracy = np.mean(predicted_labels.flatten()==test_labels.flatten())
    else:
        print('错误：输出格式有误，必须为ndarray格式')
        accuracy = 0
    return accuracy

if __name__ == '__main__':
    # 准备输出目录
    output_dir = "./step1/output/student_svm_3_less"
    os.makedirs(output_dir, exist_ok=True)
    result_path = os.path.join(output_dir, "result.md")
    figure_path = os.path.join(output_dir, "result.png")

    # 载入完整训练集
    full_train_dataset, full_train_labels = load_dataset('./step1/input/training_dataset.pkl')
    test_file = './step1/input/test_dataset_clean.pkl'

    results = []

    # 循环测试 100%, 90%, ..., 10%
    for percent in range(100, 0, -10):
        sample_size = int(full_train_dataset.shape[0] * percent / 100)
        indices = np.random.choice(full_train_dataset.shape[0], sample_size, replace=False)
        reduced_train_dataset = full_train_dataset[indices]
        reduced_train_labels = full_train_labels[indices]

        classifier = Classifier(train_dataset=reduced_train_dataset, train_labels=reduced_train_labels)
        classifier.train()
        acc = calculate_accuracy(test_file, classifier)
        results.append((percent, acc))
        print(f"使用 {percent}% 训练数据，测试集准确率: {acc:.4f}")

    # 保存结果到 markdown
    with open(result_path, "w", encoding="utf-8") as f:
        f.write("# SVM 少量训练数据实验结果\n\n")
        f.write("| 使用训练数据比例 | 测试集准确率 |\n")
        f.write("|----------------|--------------|\n")
        for percent, acc in results:
            f.write(f"| {percent}% | {acc:.4f} |\n")

    # 画图
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    percents = [p for p, _ in results]
    accuracies = [a for _, a in results]
    
    plt.figure(figsize=(10, 6))
    plt.plot(percents, accuracies, marker='o', linestyle='-', color='b', label="测试集准确率")
    plt.title("SVM在不同训练数据比例下的性能", fontsize=16)
    plt.xlabel("训练数据比例 (%)", fontsize=12)
    plt.ylabel("测试集准确率", fontsize=12)
    plt.xticks(percents)
    plt.ylim(min(accuracies) * 0.95, max(accuracies) * 1.05)
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.legend()
    plt.tight_layout()
    
    plt.savefig(figure_path)
    plt.close()

    print(f"\n结果已保存到 {result_path}")
    print(f"折线图已保存到 {figure_path}")